from math import sqrt
import torch
import torch.nn as nn
import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.losses.multiBoxLoss import MultiBoxLoss
from utils import *

class SSDHead(nn.Module):
    def __init__(self, input_channels, num_class):
        super(SSDHead, self).__init__()
        self.num_class = num_class
        self.priors_cxcy = self.create_prior_boxes()

        activation = nn.ELU()
        
        self.vgg = nn.Sequential(
            nn.Conv2d(in_channels=input_channels, out_channels=64, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
            activation,
            nn.MaxPool2d(2, ceil_mode=False),

            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
            activation,
            nn.MaxPool2d(2, ceil_mode=False),

            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
            activation,
            nn.MaxPool2d(2, ceil_mode=True),

            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            activation,
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
            activation
        )

        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, stride=2, padding=1),
            activation,
            nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=1, stride=1, padding=0),
            activation
        )

        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1, stride=1, padding=0),
            activation,
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1),
            activation
        )

        self.block3 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=128, kernel_size=1, stride=1, padding=0),
            activation,
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1),
            activation
        )

        self.block4 = nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1, stride=1, padding=0),
            activation,
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=0),
            activation
        )

        self.block5 = nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1, stride=1, padding=0),
            activation,
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=0),
            activation
        )

        self.extract1 = nn.Conv2d(in_channels=512, out_channels=4*(4+num_class), kernel_size=3, stride=1, padding=1)
        self.extract2 = nn.Conv2d(in_channels=1024, out_channels=6*(4+num_class), kernel_size=3, stride=1, padding=1)
        self.extract3 = nn.Conv2d(in_channels=512, out_channels=6*(4+num_class), kernel_size=3, stride=1, padding=1)
        self.extract4 = nn.Conv2d(in_channels=256, out_channels=6*(4+num_class), kernel_size=3, stride=1, padding=1)
        self.extract5 = nn.Conv2d(in_channels=256, out_channels=4*(4+num_class), kernel_size=3, stride=1, padding=1)
        self.extract6 = nn.Conv2d(in_channels=256, out_channels=4*(4+num_class), kernel_size=3, stride=1, padding=1)

    def create_prior_boxes(self):
        fmap_dims = {
            "conv4_3": 38,
            "conv7": 19,
            "conv8_2": 10,
            "conv9_2": 5,
            "conv10_2": 3,
            "conv11_2": 1 
        }
        obj_scales = {
            "conv4_3": 0.1 ,
            "conv7": 0.2,
            "conv8_2": 0.375,
            "conv9_2": 0.55,
            "conv10_2": 0.725,
            "conv11_2": 0.9
        }
        aspect_ratios = {
            "conv4_3": [1., 2., .5],
            "conv7": [1., 2., 3., .5, .333],
            "conv8_2": [1., 2., 3., .5, .333],
            "conv9_2": [1., 2., 3., .5, .333],
            "conv10_2": [1., 2., .5],
            "conv11_2": [1., 2., .5]
        }
        fmaps = list(fmap_dims.keys())
        
        priors_boxes = []
        
        for k, fmap in enumerate(fmap_dims):
            for i in range(fmap_dims[fmap]):
                for j in range(fmap_dims[fmap]):
                    cx = (j + 0.5) / fmap_dims[fmap]
                    cy = (i + 0.5) / fmap_dims[fmap]

                    for ratio in aspect_ratios[fmap]:
                        box_loc = [cx, cy, obj_scales[fmap]*sqrt(ratio), obj_scales[fmap]/sqrt(ratio)]
                        priors_boxes.append(box_loc)
                        
                        if ratio == 1.:
                            try:
                                additional_scale = sqrt(obj_scales[fmap] * obj_scales[fmaps[k+1]])
                            except IndexError:
                                additional_scale = 1.
                            box_loc = [cx, cy, additional_scale, additional_scale]
                            priors_boxes.append(box_loc)
        priors_boxes = torch.FloatTensor(priors_boxes).to(device)
        priors_boxes.clamp_(0, 1)
        return priors_boxes

    def ssd_decode(self, locations, num_class):
        c, h, w = locations.size()
        num_anchors = c // (4 + num_class)

        predicted_locs = []
        predicted_scores = []

        for i in range(h):
            for j in range(w):
                for ar_index in range(num_anchors):
                    offset = ar_index * (4 + num_class)
                    predicted_locs.append(locations[offset:offset + 4, i, j])
                    predicted_scores.append(locations[offset + 4:offset + 4 + num_class, i, j])

        predicted_locs = torch.stack(predicted_locs).to(device)
        predicted_scores = torch.stack(predicted_scores).to(device)

        return predicted_locs, predicted_scores

    def forward(self, x):
        feature1 = self.vgg(x)
        feature2 = self.block1(feature1)
        feature3 = self.block2(feature2)
        feature4 = self.block3(feature3)
        feature5 = self.block4(feature4)
        feature6 = self.block5(feature5)

        anchors1 = self.extract1(feature1)
        anchors2 = self.extract2(feature2)
        anchors3 = self.extract3(feature3)
        anchors4 = self.extract4(feature4)
        anchors5 = self.extract5(feature5)
        anchors6 = self.extract6(feature6)

        predicted_locs = [[] for _ in range(x.size(0))]
        predicted_scores = [[] for _ in range(x.size(0))]

        for i in range(x.size(0)):
            locs, scores = self.ssd_decode(anchors1[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

            locs, scores = self.ssd_decode(anchors2[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

            locs, scores = self.ssd_decode(anchors3[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

            locs, scores = self.ssd_decode(anchors4[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

            locs, scores = self.ssd_decode(anchors5[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

            locs, scores = self.ssd_decode(anchors6[i], self.num_class)
            predicted_locs[i].append(locs)
            predicted_scores[i].append(scores)

        predicted_locs = torch.stack([torch.cat(loc, dim=0) for loc in predicted_locs])
        predicted_scores = torch.stack([torch.cat(scores, dim=0) for scores in predicted_scores])

        return predicted_locs, predicted_scores


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SSDHead(input_channels=3, num_class=10).to(device)
input_image = torch.randn(2, 3, 300, 300).to(device)
gt_boxes = [
    torch.tensor([[0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4], [0.6, 0.6, 0.8, 0.8]]).to(device),
    torch.tensor([[0.15, 0.15, 0.25, 0.25], [0.35, 0.35, 0.45, 0.45], [0.65, 0.65, 0.85, 0.85]]).to(device)
]

# Ground truth labels
gt_labels = [
    torch.tensor([1, 2, 3], dtype=torch.long).to(device),
    torch.tensor([1, 2, 3], dtype=torch.long).to(device)
]

# Forward pass
predicted_locs, predicted_scores = model(input_image)

# Instantiate MultiBoxLoss
multibox_loss = MultiBoxLoss(model.priors_cxcy).to(device)

# Calculate loss
loss = multibox_loss(predicted_locs, predicted_scores, gt_boxes, gt_labels)

print("Loss:", loss.item())